library(data.table)
library(grid)
library(gridExtra)
library(dplyr)
library(tidyr)
library(ggplot2)
library(stringr)
library(readr)
library(latex2exp)

## Paper or presentation
paper <- TRUE

## data directory
dir.project <- './'
dir.data <- paste0(dir.project, 'data/')
dir.raw <- paste0(dir.data, 'raw/')
dir.generated <- paste0(dir.data, 'generated/')

## figures directory for output
dir.fig.paper <-
  paste0(dir.project, 'doc/paper/figures/compensation/')
dir.fig.presentation <-
  paste0(dir.project, 'results/figures/presentation/')

source(paste0(dir.project, 'src/r/figures.R'))

fileName <- paste0(dir.project, 'calibration_version.txt')
calibration_version <-
  read_file(fileName) %>% str_replace_all("[\r\n]" , "")

#############
## FIGURES ##
#############

## load data
df.optimum <-
  fread(
    paste0(
      dir.generated,
      'matlab/OptimTax/compensation_gpops/20191107_nochoice_compensation_baseline_optimum.csv'
    )
  )
df.optimum.long <-
  df.optimum %>% pivot_longer(-pricedrop) %>% mutate(version = 'optimum')

df.default <-
  fread(
    paste0(
      dir.generated,
      'matlab/OptimTax/compensation_gpops/20191107_nochoice_compensation_baseline_default_capital_tax.csv'
    )
  )
df.default.long <-
  df.default %>% pivot_longer(-pricedrop) %>% mutate(version = 'default')

df.equal <-
  fread(
    paste0(
      dir.generated,
      'matlab/OptimTax/compensation_gpops/20191107_nochoice_compensation_baseline_equipment_equal_robot_tax.csv'
    )
  )
df.equal.long <-
    df.equal %>% pivot_longer(-pricedrop) %>% mutate(version = 'equal')

df.robottax_only <-
  fread(
    paste0(
      dir.generated,
      'matlab/OptimTax/compensation_hsv/20191107_nochoice_compensation_baseline_optimal_robot_tax_hsv.csv'
    )
  )
df.robottax_only.long <-
  df.robottax_only %>% pivot_longer(-pricedrop) %>% mutate(version = 'robottax_only')



df.compensate_adjust_taus <-
  fread(
    paste0(
      dir.generated,
      'matlab/OptimTax/compensation_hsv/20191107_nochoice_compensation_baseline_optimtax_hsv_adjust_taus.csv'
    )
  )

df.compensate_adjust_taus.long <-
  df.compensate_adjust_taus %>% pivot_longer(-pricedrop) %>% mutate(version = 'capital_taxes_only')



## combine in one df
df.combined <- rbind(
  df.optimum %>% mutate(version = 'optimum'),
  df.default %>% mutate(version = 'default'),
  df.equal %>% mutate(version = 'equal'),
  df.robottax_only %>% mutate(version = 'robottax_only'),
  df.compensate_adjust_taus %>% mutate(version = 'capital_taxes_only')
) %>%
  mutate(comp_gdp_pct = compensation / Y * 100)

df.combined.long <-
  df.combined %>% pivot_longer(-c(pricedrop, version))

# 2016 USD per capita based on World Bank
gdp_per_cap <- 57588

lab <- c("Reform 1", "Reform 2", "Reform 3")

compensation_all <-
  ggplot(df.combined.long %>% filter(name == 'comp_gdp_pct', version!='robottax_only', version!='capital_taxes_only'),
         aes(x = pricedrop * 100, y = value)) + geom_line(aes(color = version, linetype =
                                                                version)) +
  ylab('pct. of GDP') + xlab('pct. drop in robot price') +
  scale_y_continuous(sec.axis = sec_axis( ~ . * gdp_per_cap / 100, name = 'USD per capita')) +
  scale_linetype_manual(
    name = '',
    values = c('solid', 'dashed', 'dotted'),
    labels = lab
  ) +
  scale_color_manual(name = '',
                     values = myPalette,
                     labels = lab) +
  theme_bw() + theme.serif +
  theme(axis.title.y.right = element_text(margin = margin(
    t = 0,
    r = 0,
    b = 0,
    l = 8
  ))) +
  theme(legend.position = c(0.3, 0.8),
        legend.background = element_rect(fill = alpha('white', 0)))


to.pdf(
  print(compensation_all),
  "compensation_all.pdf",
  width = width.double2,
  height = height.double2
)


df.comp_min_default <-
  df.combined %>% select(pricedrop, version, comp_gdp_pct) %>% filter(version %in% c('default', 'equal', 'optimum')) %>%
  pivot_wider(names_from = version, values_from = comp_gdp_pct) %>% mutate(optimum_min_default = optimum - default,
                                                                           equal_min_default = equal - default) %>%
  pivot_longer(-pricedrop)


lab2 <- c("Reform 2 - Reform 1", "Reform 3 - Reform 1")


# I use a polynomial smoothing to adjust for one numerical hickup at one price level,
# however, this only works if there are enough observations,
# that's why the if clause: with too few observations, no polynomial smoothing is applied
if (nrow(df.comp_min_default) <= 20) {
  compensation_nonlinear <-
    ggplot(df.comp_min_default %>% filter(name %in% c(
      'optimum_min_default', 'equal_min_default'
    )),
    aes(x = pricedrop * 100, y = value)) +
    geom_line(aes(color = name, linetype = name)) +
    theme_bw() + theme.serif + ylab('pct. of GDP') + xlab('pct. drop in robot price') +
    theme(axis.title.y.right = element_text(margin = margin(
      t = 0,
      r = 0,
      b = 0,
      l = 8
    ))) +
    scale_y_continuous(sec.axis = sec_axis( ~ . * gdp_per_cap / 100, name = 'USD per capita')) +
    scale_linetype_manual(
      name = '',
      values = c('dashed', 'dotted'),
      labels = lab2
    ) +
    scale_color_manual(name = '',
                       values = c("red", "blue"),
                       labels = lab2) +
    theme(
      legend.position = c(0.45, 0.9),
      legend.background = element_rect(fill = alpha('white', 0))
    )
} else {
  compensation_nonlinear <-
    ggplot(df.comp_min_default %>% filter(name %in% c(
      'optimum_min_default', 'equal_min_default'
    )),
    aes(x = pricedrop * 100, y = value)) +
    geom_smooth(
      aes(color = name, linetype = name),
      method = 'glm',
      formula = y ~ poly(x, 5),
      se = FALSE
    ) +
    theme_bw() + theme.serif + ylab('pct. of GDP') + xlab('pct. drop in robot price') +
    theme(axis.title.y.right = element_text(margin = margin(
      t = 0,
      r = 0,
      b = 0,
      l = 8
    ))) +
    scale_y_continuous(sec.axis = sec_axis( ~ . * gdp_per_cap / 100, name = 'USD per capita')) +
    scale_linetype_manual(
      name = '',
      values = c('dashed', 'dotted'),
      labels = lab2
    ) +
    scale_color_manual(name = '',
                       values = c("red", "blue"),
                       labels = lab2) +
    theme(
      legend.position = c(0.45, 0.9),
      legend.background = element_rect(fill = alpha('white', 0))
    )
}

to.pdf(
  print(compensation_nonlinear),
  "compensation_nonlinear.pdf",
  width = width.double2,
  height = height.double2
)

df.comp_differentiate <-
  df.comp_min_default %>% filter(name %in% c("optimum", "equal")) %>%
  spread(name, value) %>% mutate(diff = optimum - equal)

# plot for value of differentiating equipment and robot tax
p <-
  ggplot(df.comp_differentiate, aes(x = pricedrop * 100, y = diff)) +
  geom_line() +
  theme_bw() + ylab('pct. of GDP') + xlab('pct. drop in robot price') +
  scale_y_continuous(sec.axis = sec_axis( ~ . * gdp_per_cap / 100, name = 'USD per capita'))

to.pdf(print(p),
       "compensation_differentiation.pdf",
       width = width.double2,
       height = height.double2)



## combine plots in one
plots_combined <-
  grid.arrange(
    compensation_all + ggtitle('A: Welf. gains of reforms 1, 2, 3'),
    compensation_nonlinear + ggtitle('B: Welf. gains of adjusting capital taxes'),
    ncol = 2,
    nrow = 1
  )

ggsave(
  paste0(dir.fig, "compensation_combined.pdf"),
  plot = plots_combined,
  width = 13.5,
  height = 6.5,
  units = "cm"
)



## adjusting robot tax only
df.adjust_tau_B <-
  fread(
    paste0(
      dir.generated,
      'matlab/OptimTax/hsv/20191107_nochoice_adjust_robot_tax.csv'
    )
  )


optim_tau_B <-
  ggplot(df.adjust_tau_B,
         aes(x = pricedrop * 100, y = tau_B * 100)) + geom_line() +
  ylab('') + xlab('pct. drop in robot price') +
    ggtitle('Tax on robots in pct.') + theme_bw() + theme.serif +
    scale_y_continuous(sec.axis = sec_axis( ~ . * 1.0436/0.0436, name = ''))
  


compensation_tau_B <-
  ggplot(df.combined.long %>% filter(name == 'comp_gdp_pct', version=='robottax_only'),
         aes(x = pricedrop * 100, y = value)) + geom_line() +
  ylab('pct. of GDP') + xlab('pct. drop in robot price') +
  scale_y_continuous(sec.axis = sec_axis( ~ . * gdp_per_cap / 100, name = 'USD per capita')) +
  theme_bw() + theme.serif +
  theme(axis.title.y.right = element_text(margin = margin(
    t = 0,
    r = 0,
    b = 0,
    l = 8
  ))) +
  theme(legend.position = c(0.3, 0.8),
        legend.background = element_rect(fill = alpha('white', 0)))


plots_combined_tau_B <-
  grid.arrange(
    optim_tau_B + ggtitle('A: Tax on robots in pct.'),
    compensation_tau_B + ggtitle('B: Welfare gains'),
    ncol = 2,
    nrow = 1,
    widths=c(7.5,7.5)
  )

ggsave(
  paste0(dir.fig, "compensation_tau_B.pdf"),
  plot = plots_combined_tau_B,
  width = 13,
  height = 5,
  units = "cm"
)


## adjust all capital taxes
df.adjust_taus <-
  fread(
    paste0(
      dir.generated,
      'matlab/OptimTax/hsv/20191107_nochoice_adjust_taus.csv'
    )
  )

df.ajdust_taus.long <- df.adjust_taus %>%
    select(pricedrop, tau_B, tau_E, tau_S) %>%
    mutate(pricedrop = as.numeric(pricedrop)) %>%
    gather(starts_with('tau_'), key='tax', value='value')

optim_taus <- ggplot(df.ajdust_taus.long , aes(x = pricedrop,y=value*100)) +
    geom_line(aes(linetype=tax)) +
    theme_bw() +
    theme.serif +
    theme(legend.position = c(0.2, 0.8),
          legend.background = element_rect(fill = alpha('white', 0))) +
        xlab('pct. drop in robot price') +
    ylab('') +
    scale_y_continuous(sec.axis = sec_axis( ~ . * 1.0436/0.0436, name = ''), limits = c(-1, 30)) +
    scale_linetype_manual(
        name = '',
        values = c('tau_B' = 1, 'tau_E' = 2, 'tau_S' = 3),
        labels = unname(TeX(c(
            "$t_B$", "$t_E$", "$t_S$"
        )))
    ) 


compensation_taus <-
  ggplot(df.combined.long %>% filter(name == 'comp_gdp_pct', version=='capital_taxes_only'),
         aes(x = pricedrop * 100, y = value)) + geom_line() +
  ylab('pct. of GDP') + xlab('pct. drop in robot price') +
  scale_y_continuous(sec.axis = sec_axis( ~ . * gdp_per_cap / 100, name = 'USD per capita')) +
  theme_bw() + theme.serif +
  theme(axis.title.y.right = element_text(margin = margin(
    t = 0,
    r = 0,
    b = 0,
    l = 8
  ))) +
  theme(legend.position = c(0.3, 0.8),
        legend.background = element_rect(fill = alpha('white', 0)))


plots_combined_taus <-
  grid.arrange(
    optim_taus + ggtitle('A: Tax on capital in pct.'),
    compensation_taus + ggtitle('B: Welfare gains'),
    ncol = 2,
    nrow = 1,
    widths=c(7.5,7.5)
  )

ggsave(
  paste0(dir.fig, "compensation_taus.pdf"),
  plot = plots_combined_taus,
  width = 13,
  height = 6.5,
  units = "cm"
)

############################################
## repeat, but now in pct. of consumption ##
############################################

## in pct of consumption
dat_cons <- fread(paste0(dir.raw,'fred/A794RX0Q048SBEA.csv')) %>% rename(cons_2012 = A794RX0Q048SBEA)
cpi <- fread("./data/raw/fred/DPCERG3A086NBEA.csv")

cpi <- cpi %>% rename(date=DATE,cpi=DPCERG3A086NBEA)
cpi$year <- as.numeric(substring(cpi$date,1,4))
base16 <- cpi$cpi[cpi$year==2016]
cpi <- cpi %>% mutate(cpi16=base16/cpi)

# per cap. consumption, from 2012$ into 2016$ 
cons_per_cap <- mean(dat_cons$cons_2012) * cpi$cpi16[cpi$year==2012]

compensation_all_cons <-
  ggplot(df.combined.long %>% filter(name == 'comp_gdp_pct', version!='robottax_only', version!='capital_taxes_only'),
         aes(x = pricedrop * 100, y = value)) + geom_line(aes(color = version, linetype =
                                                                version)) +
  ylab('pct. of GDP') + xlab('pct. drop in robot price') +
  scale_y_continuous(sec.axis = sec_axis( ~ . * gdp_per_cap / cons_per_cap, name = 'pct. of cons.')) +
  scale_linetype_manual(
    name = '',
    values = c('solid', 'dashed', 'dotted'),
    labels = lab
  ) +
  scale_color_manual(name = '',
                     values = myPalette,
                     labels = lab) +
  theme_bw() + theme.serif +
  theme(axis.title.y.right = element_text(margin = margin(
    t = 0,
    r = 0,
    b = 0,
    l = 8
  ))) +
  theme(legend.position = c(0.3, 0.8),
        legend.background = element_rect(fill = alpha('white', 0)))


if (nrow(df.comp_min_default) <= 20) {
  compensation_nonlinear_cons <-
    ggplot(df.comp_min_default %>% filter(name %in% c(
      'optimum_min_default', 'equal_min_default'
    )),
    aes(x = pricedrop * 100, y = value)) +
    geom_line(aes(color = name, linetype = name)) +
    theme_bw() + theme.serif + ylab('pct. of GDP') + xlab('pct. drop in robot price') +
    theme(axis.title.y.right = element_text(margin = margin(
      t = 0,
      r = 0,
      b = 0,
      l = 8
    ))) +
    scale_y_continuous(sec.axis = sec_axis( ~ . * gdp_per_cap / cons_per_cap, name = 'pct. of cons.')) +
    scale_linetype_manual(
      name = '',
      values = c('dashed', 'dotted'),
      labels = lab2
    ) +
    scale_color_manual(name = '',
                       values = c("red", "blue"),
                       labels = lab2) +
    theme(
      legend.position = c(0.45, 0.9),
      legend.background = element_rect(fill = alpha('white', 0))
    )
} else {
  compensation_nonlinear_cons <-
    ggplot(df.comp_min_default %>% filter(name %in% c(
      'optimum_min_default', 'equal_min_default'
    )),
    aes(x = pricedrop * 100, y = value)) +
    geom_smooth(
      aes(color = name, linetype = name),
      method = 'glm',
      formula = y ~ poly(x, 5),
      se = FALSE
    ) +
    theme_bw() + theme.serif + ylab('pct. of GDP') + xlab('pct. drop in robot price') +
    theme(axis.title.y.right = element_text(margin = margin(
      t = 0,
      r = 0,
      b = 0,
      l = 8
    ))) +
    scale_y_continuous(sec.axis = sec_axis( ~ . * gdp_per_cap / cons_per_cap, name = 'pct. of cons.')) +
    scale_linetype_manual(
      name = '',
      values = c('dashed', 'dotted'),
      labels = lab2
    ) +
    scale_color_manual(name = '',
                       values = c("red", "blue"),
                       labels = lab2) +
    theme(
      legend.position = c(0.45, 0.9),
      legend.background = element_rect(fill = alpha('white', 0))
    )
}

plots_combined_cons <-
  grid.arrange(
    compensation_all_cons + ggtitle('A: Welf. gains of reforms 1, 2, 3'),
    compensation_nonlinear_cons + ggtitle('B: Welf. gains of adjusting capital taxes'),
    ncol = 2,
    nrow = 1
  )

ggsave(
  paste0(dir.fig, "compensation_combined_cons.pdf"),
  plot = plots_combined_cons,
  width = 14,
  height = 6.5,
  units = "cm"
)
